#!/usr/bin/env python
# http://en.wikipedia.org/wiki/Chord_(peer-to-peer)
# Add bootstrap method that tries to talk to the local IP - disables functionality unless/until it works
# Merge stuff from main in
# Graceful shutdown
# db entries:
#   key => ubiname
#   ubiname => host 
#   ubiname => ubimeta 
#       comes from a request for a signature
import xml.etree.ElementTree as ET
import os
import traceback

if "UBIQIS_ROOT" not in os.environ:
    os.environ["UBIQIS_ROOT"] = "/mnt/ubiqis"

import uroot.config
import ubiqis.names as names

import time
import re
import traceback
import hashlib
#import psycopg2 as dbi
from socket import getfqdn
from ubiqis import keytool
import sqlite3 as dbi
import urllib
import urllib2
import pickle
import sys
import threading
from Queue import Queue, Empty
import random
import cache
from ubiqis.util import getpkg

from getip import getip
#from socket import getfqdn
#from wsgiref import simple_server, util as wsgiref_util
import wsgiserver
import xml.etree.ElementTree as ET
from xml.parsers.expat import ExpatError

def shahex(x):
    m = hashlib.sha1()
    m.update(x)
    h = m.hexdigest()
    return int(h,16)

### Start Database Section ###

db = None #dbi.connect(database="ubiqis")
c = None #db.cursor()
debug = False

db_strf = '%s'
if dbi.__package__ == 'sqlite3':
    db_strf = '?'
"""
dbex() is designed to make a generic interface to executing
sql queries. Something that works for both sqlite3 and psycopg2
at the very least.
"""
def dbex(cursor,sql,vars=[]):
    try:
        nvars = []
        query = ''
        args = sql.split('?')
        n = len(args)
        for i in xrange(0,n):
            query += args[i]
            if i+1 == n:
                break
            if type(vars[i]) == str:
                query += db_strf
                nvars.append(vars[i])
            else:
                query += str(vars[i])
        cursor.execute(query,tuple(nvars))
    except IndexError:
        print 'Index Error: sql=',sql,' vars=',vars
        traceback.print_exc(file=sys.stdout)
        sys.exit(0)

def initdb(name):
    global db, c
    db = dbi.connect(database=name)
    c = db.cursor()
    if c == None:
        print 'could not open database'
        sys.exit(3)
    print 'Database connection opened'

    try:
        # ensure table dht is defined
        dbex(c,'select 1 from dht')
        return
    except Exception:
        pass

    dbex(c,
    '''
    create table dht (
        vkey decimal,
        tkey varchar,
        tval varchar,
        isurl varchar(1),
        primary key(tkey,tval))
    ''')
    dbex(c,'create index vkey_index on dht (vkey)')
    dbex(c,'create index tkey_index on dht (tkey)')
    dbex(c,
    '''
    create table hosts (
        vkey decimal,
        name varchar,
        state varchar,
        primary key(name)
    )
    ''')
    db.commit()
    print 'Database initialized'

def dbhosts():
    try:
        dbex(c,"select tval from dht where isurl = 'T'")
        result = c.fetchall()
        if result == None:
            return []
        else:
            return random.shuffle(c.fetchall())
    except Exception, e:
        print e
        traceback.print_exc(file=sys.stdout)
    finally:
        return []

def dball():
    try:
        dbex(c,"select tkey,tval from dht")
        return c.fetchall()
    except Exception, e:
        print e
        traceback.print_exc(file=sys.stdout)

def dbhaskey(key):
    try:
        dbex(c,"select tval from dht where tkey = ?",(key,))
        return c.fetchone() != None
    except Exception, e:
        print e
        traceback.print_exc(file=sys.stdout)

def dbinsert(key,val):
    vkey = shahex(key)
    try:
        isurl = 'F'
        if val != None and val.__class__ == str and (val.startswith("http://") or val.startswith("https://")):
            isurl = 'T'
        dbex(c,"insert into dht (vkey,tkey,tval,isurl) values (?,?,?,?)",(vkey,key,val,isurl))
    except dbi.IntegrityError:
        pass
    except Exception,e:
        print e
        traceback.print_exc(file=sys.stdout)
    db.commit()

def dbclear(key):
    dbex(c,"delete from dht where tkey = ?",(key,))
    return None

def dbget(key):
    dbex(c,"select tval from dht where tkey = ?",(key,))
    return c.fetchall()

### End Database Section ###

work_queue = Queue(1000)
big = 1<<159
MAX_HOSTS = 3
rand = random.Random()
HOST_RADIUS = 2
my_server = getip()
my_host = None
my_key = None
hosts = {}
_priv_key = keytool.key_from_file('/home/sbrandt/goog/svn/trunk/tools/registrar/private.der')
_pub_key = '/home/sbrandt/goog/svn/trunk/tools/registrar/public.pem'

def parse_args(args):
    parsed_args = {}
    for pair in args.split('&'):
        try:
            g = re.match('([^=]*)=(.*)',pair)
            if g:
                (lhs,rhs) = (g.group(1),g.group(2))
                lhs = re.sub(r'\+',' ',lhs)
                rhs = re.sub(r'\+',' ',rhs)
                parsed_args[urllib.unquote(lhs)] = urllib.unquote(rhs)
        except ValueError,e:
            print e,pair
            traceback.print_exc(file=sys.stdout)
    return parsed_args

### Start Sorting ###

def vector(a,b):
    diff = a - b
    if diff < 0 and abs(diff) > abs(diff+big):
        diff = diff+big
    elif diff > 0 and abs(diff) > abs(diff-big):
        diff = diff-big
    return diff

def distance(a,b):
    return abs(vector(a,b))

def dist_sort(hex):
    def sortk(b):
        return distance(hex,hosts[b])
    return sortk

def by_key(h):
    return hosts[h]

### End Sorting ###

### Asynchronous sending ###

"""
Do sending in an asynchronous thread so we don't block
while waiting for a response.
"""
class Sender(threading.Thread):
    def __init__(self,host,type,dict,retries=5):
        threading.Thread.__init__(self)
        if host == None or host == '':
            raise Exception("Bad Host")
        self.host = host
        self.type = type
        self.dict = dict
        self.retries = retries
    def run(self):
        url = "http://"+self.host+"/"+self.type+"?pickle="+urllib.quote(pickle.dumps(self.dict))
        for tries in xrange(0,self.retries):
            try:
                req = urllib2.Request(url)
                response = urllib2.urlopen(req)
                ans = response.read()
                break
            except urllib2.URLError, e:
                if tries == self.retries-1:
                    global work_queue
                    print 'send fail',self.host,'/',self.dict," => ",e
                    print 'fail',url
                    traceback.print_exc(file=sys.stdout)
                    work_queue.put(("/remove",{"host":self.host}))
                    break
                else:
                    time.sleep(1)

def send(host,type,dict):
    if host not in hosts:
        hosts[host] = shahex(host)
    sender = Sender(host,type,dict)
    if debug:
        print "send =>",host,type,dict
    sender.start()

### End Asynchronous sending

### Messages ###

def removehost(host):
    if host in hosts:
        del hosts[host]

def nearhosts(key,rad):
    val = shahex(key)
    hosta = hosts.keys()
    if len(hosta) == 0:
        return []
    hosta.sort(key=dist_sort(val))
    center = hosta[0]
    hosta.sort(key=by_key)
    cpos = None
    for i in xrange(0,len(hosta)):
        if hosta[i] == center:
            cpos = i
    nhosts = set()
    n = len(hosta)
    for i in xrange(0,rad):
        nhosts.add(hosta[(cpos+i)%n])
        nhosts.add(hosta[(cpos-i+n)%n])
    return list(nhosts)

def addhosts(add_hosts):
    global hosts
    added = set()
    for host in add_hosts:
        if host == None:
            pass
        elif host not in hosts:
            added.add(host)
            hosts[host] = shahex(host)
    near = nearhosts(my_host,HOST_RADIUS)
    for host in added:
        if host in near:
            send(host,"welcome",{"hosts":near})
            def getk():
                #print 'queing',host
                work_queue.put(("/sendkeys",{"host":host}))
            t = threading.Timer(10.0,getk)
            t.start()

def hello(host,checked):
    near = nearhosts(host,HOST_RADIUS)
    send_checked = set()
    addhosts(checked)
    for h in checked:
        send_checked.add(h)
    for h in near:
        send_checked.add(h)
    for h in near:
        if h == my_host:
            addhosts([host])
        elif h not in checked:
            send(h,"hello",{"host":host,"checked":send_checked})

def store(key,value,checked):
    near = nearhosts(key,HOST_RADIUS)
    #print my_host,'store',key,near,checked
    send_checked = set()
    addhosts(checked)
    for h in checked:
        send_checked.add(h)
    for h in near:
        send_checked.add(h)
    for h in near:
        if h == my_host:
            dbinsert(key,value)
            print 'stored key',key,'on',my_host
        elif h not in checked:
            send(h,"store",{"key":key,"value":value,"checked":send_checked})
            
def lookup(key,destination,checked):
    near = nearhosts(key,HOST_RADIUS)
    #print my_host,'lookup',key,near,checked
    addhosts(checked)
    addhosts([destination])
    send_checked = set()
    for h in checked:
        send_checked.add(h)
    for h in near:
        send_checked.add(h)
    for h in near:
        if h == my_host:
            if dbhaskey(key):
                send(destination,"result",{"key":key,"value":dbget(key)})
                print 'found key',key,'on',my_host
                found = True
                break
        elif h not in checked:
            send(h,"lookup",{"key":key,"destination":destination,"checked":send_checked})
            found = True
            
def sendkeys(destination,checked):
    near = nearhosts(destination,HOST_RADIUS)
    checked.add(destination)
    addhosts(checked)
    send_checked = set()
    for h in checked:
        send_checked.add(h)
    for h in near:
        send_checked.add(h)
    for h in near:
        if h == my_host:
            for (k,v) in dball():
                dbinsert(str(k),str(v))
        elif h not in checked:
            send(h,"sendkeys",{"host":destination,"checked":send_checked})

def welcome(nhosts):
    addhosts(nhosts)

def result(key,value):
    cache.put(key,value)

class Signer (threading.Thread):
    def __init__(self,input):
        threading.Thread.__init__(self)
        self.input = input
    def run(self):
        signed = getpkg(self.input, my_host, _priv_key, _pub_key)
        send(my_host,"signed",{"input":self.input,"sig":signed})
def signpkg(input):
    if dbhaskey(input): 
        dbin = dbget(input)[0][0]
        tree = ET.fromstring(dbin)
        for sig in tree.findall("signature"):
            for host in sig.findall("host"):
                if host.text == my_host:
                    cache.put(input,dbin)
                    return
    # Sign package
    signer = Signer(input)
    signer.start()

def signed(input,signed):
    try:
        tree = ET.fromstring(signed)
        if not tree.findall('error') and not tree.findall('develop'):
            # Cache signed package
            # db[input] = signed
            cache.put(input,signed)
            dbclear(input)
            dbinsert(input,signed)

        # Enable key based search
        tree = ET.fromstring(signed)
        keys = tree.findall('key')
        for key in keys:
            dbinsert(key.text,input)
    except ExpatError:
        cache.put(input,"<error>"+entity_quote(signed)+"</error>")

### End Message Processing ###

def get_arg(args,key,default):
    if key in args:
        return args[key]
    else:
        return default

class Worker (threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
    def run(self):
        my_set = set()
        my_set.add(my_host)
        while True:
            try:
                work = work_queue.get()
                (info,args) = work
                if debug:
                    print "recv =>",info,args
                if info == '/hello':
                    hello(args['host'],get_arg(args,"checked",my_set))
                elif info == '/sendkeys':
                    sendkeys(args['host'],get_arg(args,"checked",my_set))
                elif info == '/store':
                    store(args['key'],args['value'],get_arg(args,"checked",my_set))
                elif info == '/lookup':
                    lookup(args['key'],get_arg(args,'destination',my_host),get_arg(args,"checked",my_set))
                elif info == '/result':
                    result(args['key'],args['value'])
                elif info == '/remove':
                    removehost(args['host'])
                elif info == '/signed':
                    signed(args['input'],args['sig'])
                elif info == '/welcome':
                    if 'hosts' in args:
                        welcome(args['hosts'])
                elif info == '/netboot':
                    print 'netboot successful'
                    if "contacts" in args:
                        c = args["contacts"]
                        addhosts(c)
                        for contact in c:
                            send(contact,"hello",{"host":my_host,"checked":get_arg(args,"checked",my_set)})
                elif info == '/signpkg':
                    print 'signpkg=',args
                    signpkg(args['url'])
                elif info == '/initdb':
                    initdb('ubiqis-fserv-'+str(port)+'.sqlite')
            except Exception,e:
                print 'error',e
                print e.__class__
                traceback.print_exc(file=sys.stdout)

"""
Always queue up the work and return OK. We are essentially sending
asynchronous messages over http.
"""
def application(environ, start_response):
    try:
        info = environ['PATH_INFO']
        args = parse_args(environ['QUERY_STRING'])
        if "pickle" in args:
            args = pickle.loads(args["pickle"])
        work_queue.put((info,args))
        key = None
        result = ''
        if info == '/signpkg' and "url" in args:
            key = str(args['url'])
            result = cache.get(key)
            start_response("200 OK", [('Content-Type', 'text/xml')])
            return str(result)
        start_response("200 OK", [('Content-Type', 'text/html')])
        if info == '/lookup' and "key" in args and "wait" in args and args["wait"] == "yes":
            key = str(args["key"])
            result = cache.get(key)
            if "raw" in args and args["raw"] == "yes":
                data = pickle.dumps(result)
                return data
        if info == '/hello':
            result = []
            keys = hosts.keys()
            for host in keys:
                result.append([host])
            key = 'Host Query'

        # Serve files
        if info == '/file':
            print 'load file',args
            try:
                # find the base name of the file
                base = os.path.join(uroot.config.CACHE_DIR,
                    re.sub("\\.",os.path.sep,args['name']))

                # find the name of the meta data file
                meta = base + ".ubiq"

                # find the suffix
                fp = open(meta,"r")
                suffix = ET.fromstring(fp.read()).find('suffix').text
                fp.close()

                # find the name of the payload file
                source = base + "." + suffix

                # send the file
                fp = open(source,"r")

                def readblocks():
                    while True:
                        buf = fp.read(1024)
                        yield buf
                        if buf == '':
                            fp.close()
                            break

                # Want an iterable object. This supposedly chunks the output.
                return readblocks()
            except Exception,e:
                return traceback.format_exc()
        refresh = ''
        # The page should be XHTML, to enable parsing
        page = '''<html>
        <head><title>Ubiqis File Server</title>%s</head>
        <body><h1>Ubiqis File Server</h1>
        %s
        <form name='request'>
            <input type='hidden' name='key' value='%s' />
            <input type='hidden' name='val' value='%s' />
        </form>
        <hr />
        Lookup:
        <form name='lookup' action='/lookup'>
            Key <input type='text' name='key' /><br />
            <input type='submit' name='submit' value='submit' />
            <input type='hidden' name='destination' value='%s' />
            <input type='hidden' name='wait' value='yes' />
        </form>
        <hr />
        Store:
        <form name='store' action='/store'>
            Key <input type='text' name='key' /><br />
            Value <input type='text' name='value' />
            <input type='submit' name='submit' value='submit' />
        </form>
        <hr />
        Hello:
        <form name='hello' action='/hello'>
            Say, "hello", to server <input type='text' name='host' /><br />
            <input type='submit' name='submit' value='submit' />
        </form>
        </body>
    </html>'''
        qh = entity_quote(my_host)
        qr = ''
        if result == cache.NotReady:
            qr = '*** Please wait. ***'
        elif result != None:
            for r in result: 
                qr += entity_quote(str(r[0])) + "<br />"
        qk = ''
        qt = ''
        if key != None:
            qk = entity_quote(key)
            qt = "<table><tr><td>" + qk + "</td><td>=</td><td>" + qr + "</td></tr></table>"
        qp = entity_quote(pickle.dumps(result))
        return page % (refresh,qt,qk,qp,qh)
    except Exception, e:
        print 'server fail',e
        traceback.print_exc(file=sys.stdout)

def entity_quote(inp):
    dict = {'&':'&amp;','<':'&lt;','>':'&gt;','"':'&quot;',"'":"&#39;","\n":"&#10;"}
    s = ''
    for c in inp:
        if c in dict:
            s += dict[c]
        else:
            s += c
    return s

def find_offerings(arg,dirname,fnames):
    for fname in fnames:
        full = os.path.join(dirname,fname)
        if re.match(r'.*\.ubiq$',fname):
            fp = open(full,"r")
            contents = fp.read()
            tree = ET.fromstring(contents)
            suffix = tree.find('suffix').text
            name = tree.find('name').text
            lib = os.path.join(uroot.config.CACHE_DIR,
                re.sub('\.',os.path.sep,name)+"."+suffix)
            if os.path.exists(lib):
                send(my_host,"store",{"key":name,"value":my_host})
            for key in tree.findall('key'):
                send(my_host,"store",{"key":key.text,"value":name})
            fp.close()

def is_host(name):
    return re.match(r'\d+\.\d+\.\d+\.\d+:\d+',name)

def main(argv):
    global port, my_host, my_key, hosts, worker
    try:
        if len(argv) > 1:
            port = int(argv[1])
        else:
            port = 8080
        my_host = my_server + ':' + str(port)
        my_key = shahex(my_host)
        hosts[my_host] = my_key
        worker = Worker()
        work_queue.put(("/initdb",{}))
        worker.start()
        print "Running ubiqis file server at http://%s:%s" % (my_server, port)
        #httpd = simple_server.WSGIServer(('', port))
        #httpd.set_app(application)
        httpd = wsgiserver.CherryPyWSGIServer((my_server, port), application)
        if len(argv) > 2:
            contacts = set()
            for contact in argv[2:]:
                if is_host(contact):
                    contacts.add(contact)
            if len(contacts) > 0:
                send(my_host,"netboot",{"contacts":contacts})
        else:
            # If the netboot message is not received, then we
            # will send out no hellos and will not try to participate
            # in the network.
            send(my_host,"netboot",{});
        for arg in argv:
            if arg == "--walk":
                os.path.walk(uroot.config.CACHE_DIR,find_offerings,1)
        httpd.start()
        httpd.serve_forever()
    except KeyboardInterrupt:
        print "\nQuitting..."
        exit(0)

if __name__ == '__main__':
    main(sys.argv)
